'''笨鸟任务自动填充服务测试1'''

from sanic import Sanic, response
from sanic_cors import CORS, cross_origin

app = Sanic("Example")
CORS(app)

import time

import numpy as np
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)

fmt = "\n=== {:30} ===\n"
search_latency_fmt = "search latency = {:.4f}s"
num_entities, dim = 3000, 300

import json
import datetime

import loguru
import spacy
from loguru import logger

logger.info('spacy model启动加载')
time1 = datetime.datetime.now()
nlp = spacy.load('zh_core_web_md')
time2 = datetime.datetime.now()
logger.info(f'spacy model加载完成，耗时{time2 - time1}秒')

print(fmt.format("start connecting to bn_task"))
connections.connect("default", host="81.68.223.19", port="19530")
#
has = utility.has_collection("bn_task")
print(f"Does collection bn_task exist in bn_task: {has}")
#
fields = [
    FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=False),
    FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=dim)
]
schema = CollectionSchema(fields, "bn_task is the simplest demo to introduce the APIs")
print(fmt.format("Create collection `bn_task`"))
bn_task = Collection("bn_task", consistency_level="Strong")
print(fmt.format("Start loading"))
bn_task.load()


class ResponseResult:
    @staticmethod
    def sucess(data):
        return {"code": 20000, "message": '操作成功', 'data': data}


def load_json():
    with open('bn_task_index.json', 'r', encoding='utf-8') as f:
        j = json.load(f)
    # print(j)
    return dict(j)


def to_spacy(task):
    dic = {}
    token_s = nlp(task)
    emb = token_s.vector
    return emb.tolist()


def search(text):
    # print(fmt.format("Start searching based on vector similarity"))
    vectors_to_search = [to_spacy(text)]
    search_params = {
        "metric_type": "l2",
        "params": {"nprobe": 10},
    }

    start_time = time.time()
    result = bn_task.search(vectors_to_search, "embeddings", search_params, limit=10)
    end_time = time.time()

    ids = [hit.id for hit in result[0]]
    # for hits in result:
    #     for hit in hits:
    #         print(f"hit: {hit}, id: {hit.id}")
    # print(search_latency_fmt.format(end_time - start_time))
    return ids


def find_text(ids, j):
    return [{'value': j[str(id)]} for id in ids]


def test():
    j = load_json()
    ids = search('界面设计')
    result_list = find_text(ids, j)
    print(result_list)


j = load_json()


@app.route("/match")
def match(request):
    input_text = request.args['text'][0]
    ids = search(input_text)
    result_list = find_text(ids, j)
    print(result_list)
    rr = ResponseResult.sucess({
        # "parsed": True,
        # "args": request.args,
        # "url": request.url,
        "result_list": result_list
    })
    return response.json(rr)


if __name__ == '__main__':
    # test()
    ###################################################
    # Add OPTIONS handlers to any route that is missing it
    # app.register_listener(setup_options, "before_server_start")
    # Fill in CORS headers
    # app.register_middleware(add_cors_headers, "response")
    ###################################################

    app.run(host="0.0.0.0", port=8008, debug=True)
